from flask import Flask, render_template, request, render_template_string, Response
import subprocess
from prometheus_client import Counter, Histogram, generate_latest, CONTENT_TYPE_LATEST
import time

# 创建 Flask 应用实例
app = Flask(__name__)

# 创建指标
http_requests_total = Counter('http_requests_total', 'Total number of HTTP requests', ['method', 'endpoint', 'status_code'])
http_request_latency = Histogram('http_request_latency_seconds', 'HTTP request processing time', ['method', 'endpoint'])

# 存储请求开始时间
request_start_time = {}

@app.before_request
def before_request():
    # 记录请求开始时间
    request_start_time[request.url] = time.time()

@app.after_request
def after_request(response):
    # 获取请求的开始时间
    start_time = request_start_time.pop(request.url, None)
    if start_time is not None:
        # 计算请求处理时间
        duration = time.time() - start_time
 
        # 更新指标
        http_requests_total.labels(
            method=request.method,
            endpoint=request.endpoint,
            status_code=response.status_code
        ).inc()

        http_request_latency.labels(
            method=request.method,
            endpoint=request.endpoint
        ).observe(duration)

    return response

# 定义路由和视图函数
@app.route('/')
def home():
    return "Hello, Flask!"

# 动态路由示例
@app.route('/user/<username>')
def show_user(username):
    return f"Hello, {username}!"

# 表单处理示例
@app.route('/form', methods=['GET', 'POST'])
def form():
    if request.method == 'POST':
        name = request.form['name']
        return f"Hello, {name}!"
    return '''
        <form method="post">
            <input type="text" name="name" placeholder="Enter your name">
            <input type="submit" value="Submit">
        </form>
    '''

# 模板渲染示例
@app.route('/template')
def template():
    return render_template('index.html', name="World")


@app.route('/cudatest')
def cudatest():
    try:
        # 执行本地命令并捕获输出
        result = subprocess.run(
            ['/cuda-samples/vectorAdd'],
            shell=True,
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            timeout=30  # 设置超时时间（秒）
        )

        # 获取命令的输出和错误信息
        output = result.stdout
        error = result.stderr

        # 返回输出结果
        return render_template_string('''
            <h1>CUDA Test Result</h1>
            <pre>{{ output }}</pre>
            {% if error %}
                <h2>Error</h2>
                <pre>{{ error }}</pre>
            {% endif %}
        ''', output=output, error=error)

    except subprocess.TimeoutExpired:
        return "Command timed out", 500
    except subprocess.CalledProcessError as e:
        return f"Command failed with return code {e.returncode}", 500
    except Exception as e:
        return f"An error occurred: {str(e)}", 500

@app.route('/metrics')
def metrics():
    # 返回 Prometheus 指标的文本格式
    return Response(generate_latest(), mimetype=CONTENT_TYPE_LATEST)

# 启动应用
if __name__ == '__main__':
    app.run(debug=True, port=8080, host="0.0.0.0")
